import pluginModules from '../plugins/index.js';
import { getContentBetweenLast } from '../tools/str.js';
import { chatSSE } from '../service/qwen.js';
import { prisma } from '../service/mysql.js';
import { logger } from '../tools/index.js';
import { render } from '../prompts/index.js';
const JSON_REGEX = /(```json\s*([\s\S]*?)\s*```)|(`([^`]*)`)/g;

// 构建模块划分的提示词
const renderSelectModulesPrompt = (question, pluginModules) => {
    return render('modules.handlebars', {
        question, modules: pluginModules
    });
}

// 构建思考的提示词
const renderThoughtsPrompt = (question, module) => {
    return render('thoughts.handlebars', {
        role: module.role,
        examples: module.examples,
        question,
        tools: module.children
    });
}

// 构建最终回答的提示词
const renderFinalAnswerPrompt = (question, module, actions, terminal) => {
    return render('final.handlebars', {
        question, module, actions, terminal
    });
}

// 构建JSON解析的的提示词
const renderJsonParsePrompt = (inputStr) => {
    return render('json.parse.handlebars', {
        inputStr: inputStr
    });
}

// 解析思考信息
async function extractThoughtActions(text, context) {
    // 正则表达式，用于分割文本为思考块
    const thoughts = [];
    const actions = [];
    const actionInputStringList = [];
    const actionInputs = [];
    const results = [];
    let match;

    const thoughtInputPattern = /##\s*思考:\s*(.*?)(?=(\n|$))/gs;
    while ((match = thoughtInputPattern.exec(text)) !== null) {
        thoughts.push(match.slice(1)[0]);
    }

    const actionsPattern = /##\s*工具:\s*(.*?)(?=(\n|$))/gs;
    while ((match = actionsPattern.exec(text)) !== null) {
        actions.push(match.slice(1)[0]);
    }

    const actionInputsPattern = /##\s*工具输入:\s*(.*?)(?=(##\s*最终回答|$))/gs;
    while ((match = actionInputsPattern.exec(text)) !== null) {
        actionInputStringList.push(match.slice(1)[0]);
    }

    for (const actionInputString of actionInputStringList) {
        const jsonArgs = await tryExtractJson(actionInputString, context);
        actionInputs.push(jsonArgs);
    }

    actionInputs.map(o => o).forEach(actionInput => {
        const thought = thoughts.length > results.length ? thoughts[results.length] : thoughts[thoughts.length - 1];
        const action = actions.length > results.length ? actions[results.length] : actions[actions.length - 1];
        results.push({ thought, action, actionInput });
    })

    return results;
}

async function tryExtractJson(jsonStr, context) {
    try {
        return jsonParse(jsonStr)
    } catch (ignored) {
        try {
            logger.info('JSON不合法，使用大模型进行解析，inputStr = %s', jsonStr);
            const result = await chatSSE(renderJsonParsePrompt(jsonStr), (event, data) => { }, context);
            const resultText = result.output.text;
            logger.info('JSON不合法，使用大模型进行解析，result = %s', resultText);
            return jsonParse(resultText);
        } catch (err) {
            logger.error('tryExtractJson error, err = %s', err);
            return null;
        }
    }
}

function jsonParse(jsonStr) {
    try {
        return JSON.parse(jsonStr);
    } catch (ex) {
        let match;
        while ((match = JSON_REGEX.exec(jsonStr)) !== null) {
            if (match[2]) { // 匹配到了```json xxx ```
                return JSON.parse(match[2].trim());
            } else if (match[4]) { // 匹配到了`xxxx`
                return JSON.parse(match[4].trim());
            }
        }
        throw ex;
    }
}

const mergeAnswerInfo = (finalAnswer, moduleAnswer, actionProcesses = []) => {
    finalAnswer.output.finish_reason = moduleAnswer.output.finish_reason;
    finalAnswer.output.text = moduleAnswer.output.text;
    finalAnswer.usage.input_tokens = finalAnswer.usage.input_tokens += moduleAnswer.usage.input_tokens;
    finalAnswer.usage.output_tokens = finalAnswer.usage.output_tokens += moduleAnswer.usage.output_tokens;
    finalAnswer.usage.total_tokens = finalAnswer.usage.total_tokens += moduleAnswer.usage.total_tokens;
    finalAnswer.thoughts.push({
        content: moduleAnswer.output.text,
        actions: actionProcesses
    });
}

/**
 * 适配用户文本，使其更好的呈现给大模型
 */
const adapterUserText = (sourceText, context) => {
    // 替换机器人的名称为实际发言人
    sourceText = sourceText.replaceAll(`@${context.bot.name}`, ``);
    // 替换艾特人为用户标识
    for (const atUser of context.user.atWxUsers) {
        sourceText = sourceText.replaceAll(`@${atUser.name}`, `用户[${atUser.name}]`);
    }
    return sourceText;
}

// 返回代理聊天接口
export default async function (messages = [], callback = () => { }, context) {
    const userSourceQuestion = messages[messages.length - 1].content;
    const userQuestion = adapterUserText(userSourceQuestion, context);
    logger.info(`[${context.requestId}] 用户问题 = ${userQuestion}`);
    const finalAnswer = {
        prompt: userQuestion,
        output: {
            finish_reason: null,
            text: '',
        },
        thoughts: [],
        usage: {
            input_tokens: 0,
            output_tokens: 0,
            total_tokens: 0
        }
    }
    await prisma.t_message_log.create({
        data: {
            role: "user",
            content: userQuestion,
            chat_request_id: context.requestId,
            user_info: JSON.stringify(context.user)
        }
    });

    // 1. 进行模块划分
    const modulePrompt = renderSelectModulesPrompt(userQuestion, pluginModules);
    const modulePromptResult = await chatSSE(modulePrompt, () => { }, context);
    mergeAnswerInfo(finalAnswer, modulePromptResult);

    const parseModuleAnswerResult = getContentBetweenLast(modulePromptResult.output.text, '回答:', '') || modulePromptResult.output.text;
    const matchModule = pluginModules.find(module => module.name === parseModuleAnswerResult);

    logger.info(`[${context.requestId}] 模块划分 = ${parseModuleAnswerResult}`);
    if (!matchModule) {
        // 进行最终回答
        callback('end', finalAnswer);
        return finalAnswer;
    }

    // 2. 思考内容构建
    context.agentModuleName = matchModule.name;
    const plugins = matchModule.children;
    const thoughtsPrompt = renderThoughtsPrompt(userQuestion, matchModule);
    const thoughtsPromptResult = await chatSSE(thoughtsPrompt, (event, data) => {
        if (event === 'data') {
            callback('thought', data);
        }
    }, context);
    callback('end-thought', thoughtsPromptResult);

    // 3. 活动行为构建
    const actionProcesses = await extractThoughtActions(thoughtsPromptResult.output.text, context);
    logger.info(`[${context.requestId}] 行为列表 = ${JSON.stringify(actionProcesses)}`);

    // 收集信息
    mergeAnswerInfo(finalAnswer, thoughtsPromptResult, actionProcesses);

    if (actionProcesses.length === 0) {
        // 直接回答
        const finalAnswer = thoughtsPromptResult;
        const parseFinalAnswerText = getContentBetweenLast(finalAnswer.output.text, '最终回答:', '');
        if (parseFinalAnswerText) {
            finalAnswer.output.text = parseFinalAnswerText;
        }

        await prisma.t_message_log.create({
            data: {
                role: "assistant",
                content: finalAnswer.output.text,
                chat_request_id: context.requestId,
                user_info: JSON.stringify(context.user)
            }
        });
        callback('end', finalAnswer);
        return finalAnswer;
    }

    // 填充 actions 观察数据
    for (let actionProcess of actionProcesses) {
        const plugin = plugins.find(plugin => plugin.name === actionProcess.action);
        if (!plugin) {
            actionProcess.observations = actionProcess.actionInput.map((input, index) => {
                return {
                    input: input,
                    output: '工具未支持，返回为空'
                }
            });
            continue;
        }

        const resultDataList = await plugin.func(actionProcess.actionInput, context);
        actionProcess.observations = actionProcess.actionInput.map((input, index) => {
            return {
                input: input,
                output: resultDataList?.[index] || '无'
            }
        });
    }

    // 构造最终回答 Prompt
    const finalAnswerPrompt = renderFinalAnswerPrompt(userQuestion, matchModule, actionProcesses, context.terminal);
    const finalAnswerPromptResult = await chatSSE(finalAnswerPrompt, (event, data) => {
        if (event === 'data') {
            callback('answer', data);
        }
    }, context);
    finalAnswer.final_prompt = finalAnswerPrompt;
    finalAnswer.output.finish_reason = finalAnswerPromptResult.output.finish_reason;
    finalAnswer.output.text = finalAnswerPromptResult.output.text;
    finalAnswer.usage.input_tokens = finalAnswer.usage.input_tokens += finalAnswerPromptResult.usage.input_tokens;
    finalAnswer.usage.output_tokens = finalAnswer.usage.output_tokens += finalAnswerPromptResult.usage.output_tokens;
    finalAnswer.usage.total_tokens = finalAnswer.usage.total_tokens += finalAnswerPromptResult.usage.total_tokens;

    await prisma.t_message_log.create({
        data: {
            role: "assistant",
            content: finalAnswer.output.text,
            chat_request_id: context.requestId,
            user_info: JSON.stringify(context.user)
        }
    });
    callback('end', finalAnswer);
    return finalAnswer;
}